package com.aloha.middleware.db.router.strategy.impl;

import com.aloha.middleware.db.router.utils.DBContextHolder;
import com.aloha.middleware.db.router.config.DBRouterProperties;
import com.aloha.middleware.db.router.strategy.IDBRouterStrategy;

/**
 * @author DaiZhiHeng
 * @description 使用扰动计算+取模对数据进行映射
 * @date 2023/7/9 10:53
 */
public class DBRouterStrategyHashCode implements IDBRouterStrategy {

    private final DBRouterProperties properties;

    public DBRouterStrategyHashCode(DBRouterProperties properties) {
        this.properties = properties;
    }

    @Override
    public void doRouter(String dbKeyAttr) {
        int size = dbCount() * tbCount();
        int idx;
        // 扰动计算
        int hash = dbKeyAttr.hashCode() ^ (dbKeyAttr.hashCode() >>> 16);
        if ((size & (size - 1)) == 0) { // 库表的数量是 2^n
            idx = (size - 1) & hash;
        } else {
            idx = hash % size;
            idx = idx < 0 ? idx + size : idx;
        }
        // 库表索引
        int dbIdx = idx / tbCount() + 1;
        // 对于分表的名称 必须为 xxx_xxx_000 开始
        /*
            举个例子，size = 2 * 4 = 8，取模之后会得到 [0, 7] 的数
            上面的dbIdx的取值范围是 [1, 2] 即使是 dbIdx是2 idx是7 tbIdx也只到 3
            所以如果 分表是从 001 开始的 idx需要+1，或者 分表从 000 开始
         */
        int tbIdx = idx - tbCount() * (dbIdx - 1);

        // 设置到 ThreadLocal
        setDBKey(dbIdx);
        setTBKey(tbIdx);
    }

    @Override
    public void setDBKey(int dbIdx) {
        DBContextHolder.setDBKey(String.format("%02d", dbIdx));
    }

    @Override
    public void setTBKey(int tbIdx) {
        DBContextHolder.setTBKey(String.format("%03d", tbIdx));
    }

    @Override
    public int dbCount() {
        return Integer.parseInt(properties.getDbCount());
    }

    @Override
    public int tbCount() {
        return Integer.parseInt(properties.getTbCount());
    }

    @Override
    public void clear() {
        DBContextHolder.clearDBKey();
        DBContextHolder.clearTBKey();
    }
}
